# monitor/tasks.py
import os
import psutil
import gpustat
import logging
from logging import Logger


from django.utils import timezone

from monitor.models import ServerMetrics
from utils.task_util import scheduled_task


# 获取当前模块的日志记录器
logger: Logger = logging.getLogger(__name__)


@scheduled_task("performance_monitor", 60)
def performance_monitor() -> ServerMetrics:
    """收集服务器性能指标"""
    logger.info("⏳ 正在收集性能指标")
    time = timezone.now()
    # 如果在Docker中运行, 使用宿主机的系统路径
    if os.environ.get("BACKEND_IN_DOCKER") == "True":
        psutil.PROCFS_PATH = "/host/proc"  # type: ignore
    # CPU 指标
    cpu_percent = psutil.cpu_percent(interval=1)
    cpu_times = psutil.cpu_times_percent(interval=1)
    cpu_core_count = psutil.cpu_count()
    per_core_usage = psutil.cpu_percent(interval=1, percpu=True)
    # 内存指标
    virtual_memory = psutil.virtual_memory()
    # 磁盘指标
    disk_metrics = []
    for partition in psutil.disk_partitions():
        try:
            usage = psutil.disk_usage(partition.mountpoint)
            disk_metrics.append(
                {
                    "device": partition.device,
                    "mountpoint": partition.mountpoint,
                    "fstype": partition.fstype,
                    "total": usage.total,
                    "used": usage.used,
                    "free": usage.free,
                    "percent": usage.percent,
                }
            )
        except PermissionError:
            # 可能无法访问某些挂载点
            continue
    # 网络指标
    net_io = psutil.net_io_counters()
    # GPU 指标
    gpu_metrics = []
    try:
        gpu_stats = gpustat.GPUStatCollection.new_query()
        for gpu in gpu_stats.gpus:
            gpu_metrics.append(
                {
                    "index": gpu.index,
                    "uuid": gpu.uuid,
                    "name": gpu.name,
                    "temperature": gpu.temperature,
                    "fan_speed": gpu.fan_speed,
                    "memory_used": gpu.memory_used,
                    "memory_total": gpu.memory_total,
                    "gpu_utilization": gpu.utilization,
                    "power_draw": gpu.power_draw,
                    "power_limit": gpu.power_limit,
                }
            )
    except Exception as e:
        logger.error(f"❌ 获取GPU信息失败: {e}")
    # 保存到数据库
    metrics = ServerMetrics(
        time=time,
        cpu_percent=cpu_percent,
        cpu_user=getattr(cpu_times, "user", 0),
        cpu_system=getattr(cpu_times, "system", 0),
        cpu_iowait=getattr(cpu_times, "iowait", 0),
        cpu_core_count=cpu_core_count,
        per_core_usage=per_core_usage,
        memory_total=virtual_memory.total,
        memory_available=virtual_memory.available,
        memory_percent=virtual_memory.percent,
        memory_used=virtual_memory.used,
        memory_free=virtual_memory.free,
        disk_metrics=disk_metrics,
        network_bytes_sent=net_io.bytes_sent,
        network_bytes_recv=net_io.bytes_recv,
        network_packets_sent=net_io.packets_sent,
        network_packets_recv=net_io.packets_recv,
        gpu_metrics=gpu_metrics,
    )
    metrics.save()
    logger.info(f"✅ 性能指标收集完成")
    return metrics
